#!/bin/bash
#修改all_config.yaml中的mindie配置
CONFIG_FILE="./src/all_config.yaml"

eval $(awk '
BEGIN { FS = ": "; RS = "\n"; }
{
    gsub(/^[ \t]+|[ \t]+$/, "", $1);
    gsub(/^[ \t]+|[ \t]+$/, "", $2);
    if ($1 ~ /^[a-zA-Z_][a-zA-Z0-9_]*$/) {
        print toupper($1) "=" "\"" $2 "\"";
    }
}
' "$CONFIG_FILE")

ALL_IP_0_1=$(awk '
BEGIN { RS = "[\n]"; FS = ":"; }
/all_ip:/ {
    getline;
    sub(/^[ \t]+-\s*\[/, "");
    sub(/\].*$/, "");
    split($0, ips, ", ");
    for (i in ips) {
        gsub(/"/, "", ips[i]);
        if (i == 2) {
            print ips[i];
            exit;
        }
    }
}
' "$CONFIG_FILE")

source ./src/set_env_mindie_all.sh $RANK_TABLE_FILE
httpsEnabled=$HTTPSENABLED
port=$PORT
ipAddress=$ALL_IP_0_1
managementIpAddress=$ALL_IP_0_1
worldSize=$WORLDSIZE
modelWeightPath=$MODELWEIGHTPATH
modelName=$MODELNAME
multiNodesInferEnabled=$MULTINODESINFERENABLED
interNodeTLSEnabled=$INTERNODETLSENABLED

#优化配置
npuMemSize=${NPUMEMSIZE:-null}
maxSeqLen=${MAXSEQLEN:-null}
maxPrefillBatchSize=${MAXPREFILLBATCHSIZE:-null}
maxPrefillTokens=${MAXPREFILLTOKENS:-null}
maxIterTimes=${MAXITERTIMES:-null}
supportSelectBatch=${SUPPORTSELECTBATCH:-null}
maxBatchSize=${MAXBATCHSIZE:-null}
maxInputTokenLen=${MAXINPUTTOKENLEN:-null}
tokenTimeout=${TOKENTIMEOUT:-null}
e2eTimeout=${E2ETIMEOUT:-null}
tp=${TP:-null}
dp=${DP:-null}
moe_ep=${MOE_EP:-null}
moe_tp=${MOE_TP:-null}

start_device_id=$START_DEVICE_ID
npuDeviceIds=$(seq -s, $start_device_id $((start_device_id + worldSize - 1)))
echo "${ipAddress}:${port} 使用设备:${npuDeviceIds}"
jq_script='.ServerConfig.ipAddress = $ipAddress |
 .ServerConfig.managementIpAddress = $managementIpAddress |
 .BackendConfig.ModelDeployConfig.ModelConfig[0].worldSize = $worldSize |
 .BackendConfig.ModelDeployConfig.ModelConfig[0].modelWeightPath = $modelWeightPath |
 .BackendConfig.ModelDeployConfig.ModelConfig[0].modelName = $modelName |
 .BackendConfig.multiNodesInferEnabled = $multiNodesInferEnabled |
 .BackendConfig.interNodeTLSEnabled = $interNodeTLSEnabled |
 .ServerConfig.httpsEnabled = $httpsEnabled |
 .ServerConfig.port = $port |
 .ServerConfig.managementPort = $managementPort |
 .ServerConfig.metricsPort = $metricsPort |
 .BackendConfig.npuDeviceIds = [$npuDeviceIds]'

if [ "$tokenTimeout" != "null" ]; then jq_script="$jq_script | .ServerConfig.tokenTimeout = $tokenTimeout"; fi
if [ "$e2eTimeout" != "null" ]; then jq_script="$jq_script | .ServerConfig.e2eTimeout = $e2eTimeout"; fi

if [ "$maxPrefillBatchSize" != "null" ]; then jq_script="$jq_script | .BackendConfig.ScheduleConfig.maxPrefillBatchSize = $maxPrefillBatchSize"; fi
if [ "$maxPrefillTokens" != "null" ]; then jq_script="$jq_script | .BackendConfig.ScheduleConfig.maxPrefillTokens = $maxPrefillTokens"; fi
if [ "$maxBatchSize" != "null" ]; then jq_script="$jq_script | .BackendConfig.ScheduleConfig.maxBatchSize = $maxBatchSize"; fi
if [ "$maxIterTimes" != "null" ]; then jq_script="$jq_script | .BackendConfig.ScheduleConfig.maxIterTimes = $maxIterTimes"; fi
if [ "$supportSelectBatch" != "null" ]; then jq_script="$jq_script | .BackendConfig.ScheduleConfig.supportSelectBatch = $supportSelectBatch"; fi

if [ "$maxSeqLen" != "null" ]; then jq_script="$jq_script | .BackendConfig.ModelDeployConfig.maxSeqLen = $maxSeqLen"; fi
if [ "$maxInputTokenLen" != "null" ]; then jq_script="$jq_script | .BackendConfig.ModelDeployConfig.maxInputTokenLen = $maxInputTokenLen"; fi

if [ "$npuMemSize" != "null" ]; then jq_script="$jq_script | .BackendConfig.ModelDeployConfig.ModelConfig[0].npuMemSize = $npuMemSize"; fi
if [ "$tp" != "null" ]; then jq_script="$jq_script | .BackendConfig.ModelDeployConfig.ModelConfig[0].tp = $tp"; fi
if [ "$dp" != "null" ]; then jq_script="$jq_script | .BackendConfig.ModelDeployConfig.ModelConfig[0].dp = $dp"; fi
if [ "$moe_ep" != "null" ]; then jq_script="$jq_script | .BackendConfig.ModelDeployConfig.ModelConfig[0].moe_ep = $moe_ep"; fi
if [ "$moe_tp" != "null" ]; then jq_script="$jq_script | .BackendConfig.ModelDeployConfig.ModelConfig[0].moe_tp = $moe_tp"; fi

jq --arg ipAddress "$ipAddress" \
   --arg managementIpAddress "$managementIpAddress" \
   --argjson worldSize $worldSize \
   --arg modelWeightPath "$modelWeightPath" \
   --arg modelName "$modelName" \
   --argjson multiNodesInferEnabled $multiNodesInferEnabled \
   --argjson interNodeTLSEnabled $interNodeTLSEnabled \
   --argjson httpsEnabled $httpsEnabled \
   --argjson port $port \
   --argjson managementPort $(($port + 1)) \
   --argjson metricsPort $(($port + 2)) \
   --argjson npuDeviceIds [$npuDeviceIds] \
   --argjson maxSeqLen "$maxSeqLen" \
   --argjson maxPrefillBatchSize "$maxPrefillBatchSize" \
   --argjson maxPrefillTokens "$maxPrefillTokens" \
   --argjson maxIterTimes "$maxIterTimes" \
   --argjson supportSelectBatch "$supportSelectBatch" \
   --argjson maxBatchSize "$maxBatchSize" \
   --argjson maxInputTokenLen "$maxInputTokenLen" \
   --argjson tokenTimeout "$tokenTimeout" \
   --argjson e2eTimeout "$e2eTimeout" \
   --argjson npuMemSize "$npuMemSize" \
   --argjson tp "$tp" \
   --argjson dp "$dp" \
   --argjson moe_ep "$moe_ep" \
   --argjson moe_tp "$moe_tp" \
"$jq_script" /usr/local/Ascend/mindie/latest/mindie-service/conf/config.json > tmp.$$.json && mv tmp.$$.json /usr/local/Ascend/mindie/latest/mindie-service/conf/config.json

chmod 640 /usr/local/Ascend/mindie/latest/mindie-service/conf/config.json

log_dir=${LOGSDIR:-"./deepseek_logs"}
if [ ! -d "$log_dir" ]; then
    mkdir -p "$log_dir"
fi
log_file="$log_dir/${MIES_CONTAINER_IP}_mindie.log"

/usr/local/Ascend/mindie/latest/mindie-service/bin/mindieservice_daemon > "$log_file" 2>&1 &
daemon_pid=$!

while true; do
    if grep -q "Daemon start success" "$log_file"; then
        echo "${ipAddress}:${port} 服务启动成功!"
        break
    elif grep -q "kill" "$log_file"; then
        echo "检测到 'kill' 信号，退出循环。"
        exit 1
    elif ! kill -0 $daemon_pid 2>/dev/null; then
        echo "进程 ${daemon_pid} 终止。请查看日志文件: $log_file"
        exit 1
    fi
    sleep 1
done
rm -rf tmp.*.json